import sys

import pandas as pd
import matplotlib.pyplot as plt
from mpl_finance import candlestick_ohlc
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.dates import DateFormatter, MonthLocator

mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
import datetime

def get_data_no_index(code, start_day, end_day, calc_day=""):
    stock_data = pd.read_csv('d:\\csv_shares_day\\' + str(code) + '-19900101-20211231.csv', encoding="utf-8")
    stock_data.columns = ["Date", "code", "name", "close", "high", "low", "open", "pre_close", "subval",
                          "subrate",
                          "turnrate",
                          "volume", "money", "all_earn", "exchage_earn"]
    stock_data['日期'] = pd.to_datetime(stock_data['Date'])
    stock_data["开盘价"] = stock_data["open"]
    stock_data["最高价"] = stock_data["high"]
    stock_data["最低价"] = stock_data["low"]
    stock_data["收盘价"] = stock_data["close"]
    stock_data["成交量"] = stock_data["volume"]
    stock_data["成交额"] = stock_data["money"]
    stock_data.set_index('日期', inplace=True)
    return stock_data

code = "000001"
data = get_data_no_index(code, "2021-01-01", "2021-12-01")

import backtrader as bt
class DoubleMovingAverageCross(bt.Strategy):
    params = (
        ('short_window', 5),
        ('long_window', 20),
    )

    def __init__(self):
        # 计算短期和长期移动平均线
        self.sma_short = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.short_window)
        self.sma_long = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.long_window)
        # 交叉信号
        self.crossover = bt.indicators.CrossOver(self.sma_short, self.sma_long)

    def next(self):
        if not self.position:
            # 短期均线上穿长期均线，买入
            if self.crossover > 0:
                self.buy()
        else:
            # 短期均线下穿长期均线，卖出
            if self.crossover < 0:
                self.sell()

# 创建 Cerebro 引擎
cerebro = bt.Cerebro()

# 添加数据到 Cerebro
data_feed = bt.feeds.PandasData(dataname=data)
cerebro.adddata(data_feed)

# 添加策略
cerebro.addstrategy(DoubleMovingAverageCross)

# 设置初始资金
cerebro.broker.setcash(100000.0)

# 设置交易佣金，假设为 0.1%
cerebro.broker.setcommission(commission=0.001)

# 打印初始资金
print(f'初始资金: {cerebro.broker.getvalue():.2f}')

# 运行回测
cerebro.run()

# 打印最终资金
print(f'最终资金: {cerebro.broker.getvalue():.2f}')

# 绘制结果
cerebro.plot()

# 计算收益率
initial_value = 100000.0
final_value = cerebro.broker.getvalue()
return_rate = (final_value - initial_value) / initial_value
print(f'收益率: {return_rate * 100:.2f}%')